Latent neural process (LNP): Derivation of an ELBO-like objective

Latent neural processes (LNPs) #1 use a training objective that is inspired from the ELBO. The steps for the derivation are the following:

• First, we derive an ELBO that bounds the conditional marginal likelihood $p(y_T\mid x_D,y_C)$ instead of the usual $p(y_D\mid x_D)$. This marginal likelihood is more representative of the desired task in NPs, i.e. recover the behaviour of the function in unseen (target, $T$) regions given known (context, $C$) regions.
• Second, approximate the previous ELBO, which is intractable.

Suppose we have a dataset $D=\{(x_i,y_i{)}_{i=1}^n\}$ that has been generated from a latent variable $z$ by a model $p_{\theta }(y_D\mid z,x_D)$. We could estimate the parameters $\theta$ by maximum likelihood, by maximizing the marginal likelihood of $y_D={y_i}_{i=1}^n$. The log marginal likelihood can be decomposed as

which results in the usual ELBO bound

Note that this is equivalent to the ELBO bound of a VAE for a single datapoint $\mathbf{y}:=y_D$ originating from a latent variable $z$, with the only difference that here we are conditioning everything on the inputs $x_D$ too (by design).

Now, suppose that you partition the dataset into a context set $C=\{(x_i,y_i{)}_{i=1}^m\}$ and a target set $T=\{(x_i,y_i{)}_{i=m+1}^n\}$. The goal is to infer $y_T$ given information from $y_C$. We could try to obtain appropriate parameters $\theta$ by maximizing the same marginal likelihood as before, $p(y_D\mid x_D)$. However, this is an indirect objective, since it represents maximizing the likelihood of the entire dataset. What we really want to maximize is the conditional marginal likelihood $p(y_T\mid x_D,y_C)$. Following the same VAE analogy as before, this would be equivalent to reconstructing part of a datapoint based on the rest of the datapoint (for example, reconstructing the left side of an image based on the right side).

We can obtain an ELBO bound for this conditional marginal likelihood as follows. The LHS of the usual ELBO can be rewritten as

Similarly, the RHS can be rewritten as

We can obtain the desired ELBO by substracting $\log p(y_C\mid x_D)$ from both sides and $\text{collecting the operands in red}$ in the RHS. In this way, after substracting this term the LHS becomes the desired bound. The RHS becomes

where the last term cancels out:

Thus, we arrive at

which becomes the desired bound after making the (reasonable) assumption that each $y_i$ prediction can only be informed by its corresponding $x_i$ input, for example so that $p(y_T\mid z,x_D)$ becomes $p(y_T\mid z,x_T)$:

We would be done except for one problem: this expression is unfortunately intractable because $p(z\mid x_C,y_C)=\frac{p(y_C\mid x_C,z)p(z)}{\int p(y_C\mid x_C,z)p(z)\mathrm{d}y}$ is intractable.

Note that, since we have substracted the same term from the marginal likelihood and the ELBO, the overall KL divergence remains the same with respect to the original one, i.e. the overall KL divergence is still

The previous ELBO is a bound of the desired marginal likelihood $p(y_T\mid x_D,y_C)$, but it is intractable because $p(z\mid x_C,y_C)$. LNPs circumvent this issue by approximating this term with $q(z\mid x_C,y_C)$.

The right term (the KL or regularization term) can now be interpreted in the following way: the variational distribution over the latent variable $z$ should be the same when the model has access to full information about the function ($x_D,y_D$) and when the model has only partial information about the function ($x_C,y_C$). This seems a reasonable objective for NPs, which try to recover the whole $y_D,x_D$ based only on $y_C,x_C$.

Note, however, that this ELBO-like objective is no longer an analytical lower bound of the conditional log marginal likelihood $p(y_T\mid x_D,y_C)$, so there is no guarantee that we are maximizing the likelihood of the parameters anymore.

1 Garnelo et al 2018. Neural processes.

2 Garnelo et al 2018. Conditional neural processes.